package main

import (
	"encoding/binary"
	//	"fmt"
	"io"
	"net"
	"strconv"
	"time"
)

//normal connection
func NormalConn(conn net.Conn, head []byte) {
	connlocal := conn
	//defer connlocal.Close();
	//define normal target ip and port
	//fmt.Printf("get Head:%s\n", string(head))
	//fmt.Println(head)
	target := "127.0.0.1:22"
	targetConn, err := net.DialTimeout("tcp", target, 30*time.Second)
	connlocal.SetReadDeadline(time.Now().Add(time.Minute * 3))
	targetConn.SetReadDeadline(time.Now().Add(time.Minute * 3))
	if err != nil {
		//fmt.Printf("connect %s fail\n", target)
		return

	} else {
		//fmt.Printf("connect %s success\n", target)

		targetConn.Write(head)
		go func() {
			_, err = io.Copy(targetConn, connlocal)
			if err != nil {
				//fmt.Println("io.Copy 1 failed:", err.Error())
			}
		}()
		go func() {
			_, err = io.Copy(connlocal, targetConn)
			if err != nil {
				//fmt.Println("io.Copy 1 failed:", err.Error())
			}
		}()

	}
}

// port forward
func HandleConn(conn net.Conn, targetAddr string) {
	//defer conn.Close()
	targetConn, err := net.DialTimeout("tcp", targetAddr, 30*time.Second)
	targetConn.SetDeadline(time.Now().Add(time.Minute * 3))
	conn.SetDeadline(time.Now().Add(time.Minute * 3))
	if err != nil {
		//fmt.Printf("connect %s fail:%s\n", targetAddr)
		return
	} else {
		go func() {
			_, err = io.Copy(targetConn, conn)
			if err != nil {
				//fmt.Println("io.Copy 1 failed:", err.Error())
			}
		}()
		go func() {
			_, err = io.Copy(conn, targetConn)
			if err != nil {
				//fmt.Println("io.Copy 1 failed:", err.Error())
			}
		}()
	}
}

// check the head,parse ip and port
func checkHead(head []byte) (targetAddr string, err error) {
	//fmt.Printf("checkHead head:%v\n", head)
	ipstr := net.IPv4(head[0], head[1], head[2], head[3]).String()
	//fmt.Printf("get ipstr from head:%s\n", ipstr)
	if ipstr == "" {
		//fmt.Printf("invalid ip\n")
		return
	}
	portstr := strconv.Itoa(int(binary.BigEndian.Uint16(head[4:6])))
	//fmt.Printf("get portstr from head:%s\n", portstr)
	if portstr == "" {
		//fmt.Printf("invalid port\n")
		return
	}
	tokenstr := string(head[6:38])
	//fmt.Printf("get tokenstr from head:%s\n", tokenstr)
	if tokenstr != "b60d84d4168a2bf7ced43dcf29e7b18e" {
		//fmt.Printf("invalid token\n")
		return
	}
	//fmt.Printf("checkHead:%s:%s--%s\n", ipstr, portstr, tokenstr)
	return ipstr + ":" + portstr, nil

}

type MarkHead struct {
	dstIP   net.IP
	dstPort uint16
	token   string
}

func main() {
	listen, err := net.Listen("tcp", ":8080")
	if err != nil {
		//fmt.Println("listen error:", err)
		return
	}
	//fmt.Printf("listen success on :8080\n")
	for {
		//fmt.Printf("wait connections...\n")
		conn, err := listen.Accept()
		if err != nil {
			//fmt.Println("accept error:", err)
			break
		}
		//fmt.Printf("accepts one connect:%s\n", conn.RemoteAddr().String())
		head := make([]byte, 38)
		n, err := conn.Read(head)
		//fmt.Printf("Get Head %d bytes in main:%v\n", n, head)
		targetAddr, err1 := checkHead(head)
		if n != 38 || err != nil || err1 != nil {
			//fmt.Printf("run in normal conn\n")
			go NormalConn(conn, head[:n])
		} else {
			//fmt.Printf("run in special conn\n")
			go HandleConn(conn, targetAddr)
		}
	}

}
